# Copyright 2022 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for perfkitbenchmarker.time_triggers.maintenance_simulation_trigger."""

import time
import unittest
from unittest import mock
from absl import flags
from perfkitbenchmarker import benchmark_spec
from perfkitbenchmarker import sample
from perfkitbenchmarker import virtual_machine
from perfkitbenchmarker.providers.gcp import gce_virtual_machine
from tests import pkb_common_test_case
from perfkitbenchmarker.time_triggers import base_disruption_trigger
from perfkitbenchmarker.time_triggers import maintenance_simulation_trigger

FLAGS = flags.FLAGS


class MaintenanceSimulationTest(pkb_common_test_case.PkbCommonTestCase):

  def setUp(self):
    super().setUp()
    FLAGS.simulate_maintenance = True

  def testInitialization(self):
    FLAGS.simulate_maintenance_delay = 10
    FLAGS.capture_live_migration_timestamps = True
    trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
    self.assertTrue(trigger.capture_live_migration_timestamps, True)
    self.assertEqual(trigger.delay, 10)

  def testTrigger(self):
    vm = mock.create_autospec(gce_virtual_machine.GceVirtualMachine)
    trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
    trigger.vms = [vm]
    trigger.SetUp()
    self.assertIn(vm, trigger.gce_simulate_maintenance_helpers)

  @mock.patch(
      'perfkitbenchmarker.time_triggers.maintenance_simulation_trigger.GCESimulateMaintenanceFactory'
  )
  def testSetup(self, mock_factory):
    mock_helper = mock.create_autospec(
        maintenance_simulation_trigger.GCESimulateMaintenanceTool
    )
    mock_factory.return_value = mock_helper
    vm = mock.create_autospec(gce_virtual_machine.GceVirtualMachine)
    trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
    trigger.vms = [vm]
    trigger.capture_live_migration_timestamps = True
    trigger.SetUp()
    mock_helper.SetupLMNotification.assert_called_once()
    trigger.TriggerMethod(vm)
    mock_helper.StartLMNotification.assert_called_once()

  @mock.patch.object(time, 'time', mock.MagicMock(return_value=0))
  def testWaitForDisruption(self):
    vm = mock.create_autospec(virtual_machine.BaseVirtualMachine)
    vm_spec = mock.MagicMock(spec=benchmark_spec.BenchmarkSpec)
    event = base_disruption_trigger.DisruptionEvent(
        start_time=0.0, end_time=0.0, total_time=10.0
    )
    s = []
    trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
    self.enter_context(
        mock.patch.object(trigger, 'WaitForDisruption', return_value=[event])
    )
    trigger.capture_live_migration_timestamps = True
    trigger.vms = [vm]
    trigger.AppendSamples(None, vm_spec, s)
    self.assertEqual(
        s,
        [
            sample.Sample(
                'LM Total Time',
                10,
                'seconds',
                {
                    'LM_total_time': 10.0,
                    'Host_maintenance_end': 0.0,
                    'Host_maintenance_start': 0.0,
                },
                timestamp=0,
            )
        ],
    )

  def testWaitForDisruptionReturnsCorrectValue(self):
    vm = mock.create_autospec(gce_virtual_machine.GceVirtualMachine)
    vm.instance_number = 0
    trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
    trigger.gce_simulate_maintenance_helpers[vm] = mock.create_autospec(
        maintenance_simulation_trigger.GCESimulateMaintenanceTool
    )
    event = base_disruption_trigger.DisruptionEvent(
        start_time=0.0, end_time=0.0, total_time=10.0
    )
    trigger.gce_simulate_maintenance_helpers[
        vm
    ].CollectLMNotificationsTime.return_value = event
    trigger.capture_live_migration_timestamps = True
    trigger.vms = [vm]
    self.assertEqual(
        trigger.WaitForDisruption(),
        [event],
    )


if __name__ == '__main__':
  unittest.main()
